#include <fstream>
#include <string>
#include <stdio.h>
#include <iostream>
#include <complex>
#include <cmath>

#include "TFile.h"
#include "TTree.h"
#include "TBranch.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TGraph.h"
#include "TGraph2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TMath.h"
#include "TStyle.h"
#include "TCanvas.h"

#include "ToDetection.h"

using namespace std;

const int Rdim = 768;
const int Tdim = 1024;
const int TrkNo = 4;			// tracker No.
const int XYdim = 2;			// X/Y 
const double factor = 1;
const double rm = 200/factor;             // curvature of the spherical mirror, unit: cm
const double HitCut = 2.0;      // unit: cm

//parameters to solve the quartic function
const double PI = 3.1415926535897932384626;
const double M_2PI = 2*PI;
const double eps = 1e-12;

const double Xdir = -0.0;   //
const double Ydir = -0.0;   //
const double Zdir = 1;
const double Xinj = 0./factor;    // if factor = 1, unit: cm  //
const double Yinj = 0./factor;    //
const double Zinj = 0/factor;

ToDetection::ToDetection(string Path)
{
    TFile * f = new TFile(Form("./data/Photon_%s.root", Path.c_str()));
    TFile * fOut = new TFile(Form("./result/Out_%s.root", Path.c_str()), "RECREATE");

    TTree * Photon = (TTree* )f -> Get("OpticalPhoton");
    f -> GetObject("OpticalPhoton", Photon);

    Initialization();

    int Event;
    double Omit_x;
    double Omit_y;
    double Omit_z;
    double dirX;
    double dirY;
    double dirZ;
    double Final_x;
    double Final_y;
    double Final_z;
    double KEnergy;
    double ParentID;

    TBranch * b_Event;
    TBranch * b_Omit_x;
    TBranch * b_Omit_y;
    TBranch * b_Omit_z;
    TBranch * b_Dir_X;
    TBranch * b_Dir_Y;
    TBranch * b_Dir_Z;
    TBranch * b_Final_x;
    TBranch * b_Final_y;
    TBranch * b_Final_z;
    TBranch * b_KEnergy;
    TBranch * b_ParentID;

    Photon -> SetBranchAddress("Event",    &Event,    &b_Event);
    Photon -> SetBranchAddress("Omit_x",   &Omit_x,   &b_Omit_x);
    Photon -> SetBranchAddress("Omit_y",   &Omit_y,   &b_Omit_y);
    Photon -> SetBranchAddress("Omit_z",   &Omit_z,   &b_Omit_z);
    Photon -> SetBranchAddress("Dir_X",    &dirX,     &b_Dir_X);
    Photon -> SetBranchAddress("Dir_Y",    &dirY,     &b_Dir_Y);
    Photon -> SetBranchAddress("Dir_Z",    &dirZ,     &b_Dir_Z);
    Photon -> SetBranchAddress("Final_x",  &Final_x,  &b_Final_x);
    Photon -> SetBranchAddress("Final_y",  &Final_y,  &b_Final_y);
    Photon -> SetBranchAddress("Final_z",  &Final_z,  &b_Final_z);
    Photon -> SetBranchAddress("KEnergy",  &KEnergy,  &b_KEnergy);
    Photon -> SetBranchAddress("ParentID", &ParentID, &b_ParentID);

    Long64_t nentries;
    nentries = Photon -> GetEntriesFast();
    cout << "Entries = " << nentries << endl;
    int NEvent = Photon -> GetEntry(Photon -> LoadTree(nentries-1));

    for (int m = 0; m < NEvent; m ++)
    {
        vector<pair<double, double>> HitsPosition;
        vector<double> EmitPars;    // emit x, y, z; dir x, y, z
        vector<pair<double, double>> DetectionPoint;
        for (int i = 0; i < nentries; i++)
        {
            Long64_t ii = Photon -> LoadTree(i);
            Photon -> GetEntry(ii);

            if (Final_x <= -500 && Final_x >= -513)
                if (Event == m) {/*if ( m == 1) cout << "Y = " << Final_y/10 << "  Z = " << Final_z/10 << endl; */HitsPosition.push_back(make_pair(Final_y/10, Final_z/10)); }

            if (i < 100 && m == 0)
            {
                EmitPars.push_back((Omit_x+1500)/10);
                EmitPars.push_back(Omit_y);
                EmitPars.push_back(Omit_z);
                EmitPars.push_back(dirX);
                EmitPars.push_back(dirY);
                EmitPars.push_back(dirZ);

                double * CD = new double[2]; 
                CalculateDetection(EmitPars);
                EmitPars.clear();
                DetectionPoint.clear();
                double distance = sqrt(pow((CD[0] - Final_y/10),2) + pow((CD[1] - Final_z/10),2));
                cout << "distance = " << distance << endl;
                h1_Distance -> Fill(distance);
            }
        }

        vector<double> AC = AngleCalculation(HitsPosition);
        HitsPosition.clear();

        //for (int i = 0; i < AC.size(); i++) {h1_Angle -> Fill(AC[i]);}
        AC.clear();
    }
    
    h1_Angle -> Fit("gaus", "", "", 10, 40);

    h1_Angle -> Write();
    h1_RICHRadius -> Write();
    h2_RICHHitPos -> Write();
    h1_ReflectionAngle ->Write();
    h1_Theta -> Write();
    h1_Distance -> Write();

    f -> Close();
    fOut -> Close();
}

ToDetection::~ToDetection()
{}

void ToDetection::Initialization()
{
    cout << "Initialize Histograms .. " << endl;

    h1_RICHRadius = new TH1D("RICHRadius", "RICH Radius", 240, 0, 20);				        SetHG(h1_RICHRadius, "Radius/cm");
    h2_RICHHitPos = new TH2D("RICHHitPosition", "RICH Hit Position", 600, -10, 10, 600, -10, 10);		    SetHG(h2_RICHHitPos, "X/cm", "Y/cm");
    h1_RICHQSum   = new TH1D("RICHQSum", "RICH Charge Sum", 5000, 0, 50000);			    SetHG(h1_RICHQSum, "ADC");
    h1_RICHHits   = new TH1D("RICHHits", "RICH Hit Number", 10, 0, 10);				        SetHG(h1_RICHHits, "Hit Number");
    h1_Angle      = new TH1D("RICHAngle", "Cherenkov Angle Distribution", 1000, 0, 100);	SetHG(h1_Angle, "Angle/mrad");
    h1_ReflectionAngle = new TH1D("ReflectionAngle", "Reflection Angle", 200, 0, 100);      SetHG(h1_ReflectionAngle, "Angle/rad");
    h1_Theta      = new TH1D("Theta", "Polar Angle", 100, 0, 100);                          SetHG(h1_Theta, "Angle/mrad");
    h1_Distance   = new TH1D("Distance", "Distance", 100, 0, 100);                          SetHG(h1_Distance, "/cm");
}

double ToDetection::DichotomyRecursion(double Omega, double re, double r, double rm)
{
    // dichtomy recursion to solve the angular function
    double high = 200/1000;
    double low = 0/1000;
    double angle = 100;
    angle /= 1000;     // initial value is 100 mrad
    double Omegae = angle;
    double func;
    func = Omegae + asin(re*sin(Omegae)/r) - 2*asin(re*sin(Omegae)/rm) - Omega;
    //cout << "Initialization : func = " << func << "  Omegae = " << Omegae << "  Omega = " << Omega << "  re = " << re << "  r = " << r << "  rm = " << rm << endl;
    while (abs(func) > 1e-10)
    {
        if (func > 0)      high = Omegae; 
        else if (func < 0) low = Omegae;
        Omegae = (low + high)/2; 
        func = Omegae + asin(re*sin(Omegae)/r) - 2*asin(re*sin(Omegae)/rm) - Omega;

        if (low == high ) break;
        //cout << "High = " << high << "    low = " << low << endl;
    }

    return Omegae;
}

vector<double> ToDetection::AngleCalculation(vector<pair<double, double>> HitsPosition)
{
    vector<double> Angle;
    // in the coordinates transform ,the X/Y are the same in the Cosmic Ray test coordiante and center-of-curvature coordinate; only the Z axis is different
    // the hit position is the ionization point in the photon detector, obtained by extrapolating the cosmic rat tracking
    double RHitX = (Xinj + 15*Xdir/Zdir)/factor, RHitY = (Yinj + 15*Ydir/Zdir)/factor;            // unit: cm, the Cosmic Ray test coordinates and the center-of-curvature coordinates
    double RHitZ = 100/factor;                                                                 // unit: cm, the center-of-curvature coordinates
    double EmitX = 50*Xdir/Zdir/factor, EmitY = 50*Ydir/Zdir/factor;    // unit: cm, the Cosmic Ray test coordinates and the center-of-curvature coordinates
    double EmitZ = 150/factor;                                                                 // unit: cm, the center-of-curvature coordinates
    //cout << "First EmitX/Y/Z = " << EmitX << "  " << EmitY << "  " << EmitZ << endl;
    double r = 0;   // the distance between center-of-curvature and detection point
    double re = 0;  // the distance between center-of-curvature and photon emission point 
    double Cosomega = 0;        // angle omega is that of the photon emission tracing and photon detection tracing to the center-of-curvature
    double Sinomega = 0;
    double Omegae = 0, Omegaa = 0, Thetam = 0;

    for (int i = 0; i < HitsPosition.size(); i++)
    {
        double X = HitsPosition[i].first, Y = HitsPosition[i].second;
        //cout << "X = " << X << "  Y = " << Y << endl;
        h1_RICHRadius -> Fill(sqrt(pow(X-RHitX,2) + pow(Y-RHitY,2)));

        // emit tracing
        double DeteVecX = X/factor, DeteVecY = Y/factor, DeteVecZ = 100/factor;              // unit: cm, detection point
        r = sqrt(pow(DeteVecX,2) + pow(DeteVecY,2) + pow(DeteVecZ,2));
        double EmitVecX = EmitX, EmitVecY = EmitY, EmitVecZ = (r+rm)/2;    // unit: cm, emit point
        //re = sqrt(pow(EmitX,2) + pow(EmitY,2) + pow(EmitZ,2)); 
        re = sqrt(pow(EmitVecX,2) + pow(EmitVecY,2) + pow(EmitVecZ,2));
        Cosomega = (EmitVecX*DeteVecX + EmitVecY*DeteVecY + EmitVecZ*DeteVecZ)/re/r;
        if ((EmitVecX*DeteVecX + EmitVecY*DeteVecY) < 0) Sinomega = -sqrt(1 - pow(Cosomega,2));
        else Sinomega = sqrt(1 - pow(Cosomega,2));

        double Thetap = 0, Phip = 0;
        Thetap = acos(Zdir/sqrt(pow(Xdir,2) + pow(Ydir,2) + pow(Zdir,2)));

        if (Xdir != 0) 
        {
            if (Xdir < 0) 
	        {
		        if (Ydir > 0) Phip = acos(Xdir/sqrt(pow(Xdir, 2) + pow(Ydir,2)));
		        else if (Ydir < 0) Phip = 2*PI - acos(Xdir/sqrt(pow(Xdir, 2) + pow(Ydir,2)));
                else if (Ydir == 0) Phip = PI;
	        }
	        else if (Xdir > 0)
	        {
		        if (Ydir > 0) Phip = acos(Xdir/sqrt(pow(Xdir, 2) + pow(Ydir,2)));
		        else if (Ydir < 0) Phip = 2*PI - acos(Xdir/sqrt(pow(Xdir, 2) + pow(Ydir,2)));
                else if (Ydir == 0) Phip = 0;
	        }
        }
        else 
        {
            if (Ydir > 0) Phip = PI/2;
            else Phip = 3*PI/2;
        }

        //cout << "Thetap = " << Thetap << "  Phip = " << Phip << endl;

        double apz = cos(Thetap), apx = sin(Thetap)*cos(Phip), apy = sin(Thetap)*sin(Phip);
        double bpz = -apx, bpx = pow(sin(Phip),2) + cos(Thetap)*pow(cos(Phip),2), bpy = (cos(Thetap)-1)*cos(Phip)*sin(Phip);
        double cpz = -apy, cpx = bpy, cpy = pow(cos(Phip),2) + cos(Thetap)*pow(sin(Phip),2);

        double DR = DichotomyRecursion(acos(Cosomega), re, r, rm);
             
        Omegae = DR;
        double u = DeteVecZ*apz+DeteVecX*apx+DeteVecY*apy;
        double ue = EmitVecZ*apz+EmitVecX*apx+EmitVecY*apy;
        double Costheta = sin(Omegae)/Sinomega*u/r - sin(Omegae-acos(Cosomega))/Sinomega*ue/re;
        //cout << "      Cherenkov Angle = " << acos(Costheta)*1000 << endl;
        Angle.push_back(acos(Costheta)*1000);
        h2_RICHHitPos -> Fill(X, Y);

    }

    return Angle;    
}

void ToDetection::SetHG(TGraph * g)
{
    g -> SetLineWidth(2);
    g -> SetMarkerStyle(20);
    g -> SetMarkerColor(2);
    g -> SetMarkerSize(2);
    g -> GetXaxis() -> SetTitle("Z/cm");
    g -> GetXaxis() -> SetLimits(-8, 8);
    g -> GetXaxis() -> SetLabelSize(0.043);
    g -> GetXaxis() -> SetTitleSize(0.045);
    g -> GetYaxis() -> SetRangeUser(0, 150);
    g -> GetYaxis() -> SetTitle("XY/cm");
    g -> GetYaxis() -> SetLabelSize(0.043);
    g -> GetYaxis() -> SetTitleSize(0.045);
    g -> Draw("AP*");
}

void ToDetection::SetHG(TH1D *h1, string title)      // set the histogram/graph axis, line ... 
{
    h1 -> SetLineWidth(3);
    h1 -> GetXaxis() -> SetTitle(title.c_str());
    h1 -> GetXaxis() -> SetLabelSize(0.043);
    h1 -> GetXaxis() -> SetTitleSize(0.045);
    h1 -> GetYaxis() -> SetTitle("Counts");
    h1 -> GetYaxis() -> SetLabelSize(0.043);
    h1 -> GetYaxis() -> SetTitleSize(0.045);
}

void ToDetection::SetHG(TH1D *h1, string Xtitle, string Ytitle)      // set the histogram/graph axis, line ... 
{
    h1 -> SetLineWidth(3);
    h1 -> GetXaxis() -> SetTitle(Xtitle.c_str());
    h1 -> GetXaxis() -> SetLabelSize(0.043);
    h1 -> GetXaxis() -> SetTitleSize(0.045);
    h1 -> GetYaxis() -> SetTitle(Ytitle.c_str());
    h1 -> GetYaxis() -> SetLabelSize(0.043);
    h1 -> GetYaxis() -> SetTitleSize(0.045);
}

void ToDetection::SetHG(TH2D *h2, string Xtitle, string Ytitle)      // set the histogram/graph axis, line ... 
{
    h2 -> SetLineWidth(3);
    h2 -> GetXaxis() -> SetTitle(Xtitle.c_str());
    h2 -> GetXaxis() -> SetLabelSize(0.043);
    h2 -> GetXaxis() -> SetTitleSize(0.045);
    h2 -> GetYaxis() -> SetTitle(Ytitle.c_str());
    h2 -> GetYaxis() -> SetLabelSize(0.043);
    h2 -> GetYaxis() -> SetTitleSize(0.045);
}

vector<pair<double, double>> ToDetection::CalQuartic(double a, double b, double c, double d)
{
	double a3 = -b;
	double b3 =  a*c -4.*d;
	double c3 = -a*a*d - c*c + 4.*b*d;

	// cubic resolvent
	// y^3 − b*y^2 + (ac−4d)*y − a^2*d−c^2+4*b*d = 0

	double x3[3];
	unsigned int iZeroes = CalQubic(x3, a3, b3, c3);

	double q1, q2, p1, p2, D, sqD, y;

	y = x3[0];
	// THE ESSENCE - choosing Y with maximal absolute value !
	if(iZeroes != 1)
	{
		if(fabs(x3[1]) > fabs(y)) y = x3[1];
		if(fabs(x3[2]) > fabs(y)) y = x3[2];
	}

	// h1+h2 = y && h1*h2 = d  <=>  h^2 -y*h + d = 0    (h === q)

	D = y*y - 4*d;
	if(fabs(D) < eps) //in other words - D==0
	{
		q1 = q2 = y * 0.5;
		// g1+g2 = a && g1+g2 = b-y   <=>   g^2 - a*g + b-y = 0    (p === g)
		D = a*a - 4*(b-y);
		if(fabs(D) < eps) //in other words - D==0
			p1 = p2 = a * 0.5;

		else
		{
			sqD = sqrt(D);
			p1 = (a + sqD) * 0.5;
			p2 = (a - sqD) * 0.5;
		}
	}
	else
	{
		sqD = sqrt(D);
		q1 = (y + sqD) * 0.5;
		q2 = (y - sqD) * 0.5;
		// g1+g2 = a && g1*h2 + g2*h1 = c       ( && g === p )  Krammer
		p1 = (a*q1-c)/(q1-q2);
		p2 = (c-a*q2)/(q1-q2);
	}

    complex<double> retval[4];

	// solving quadratic eq. - x^2 + p1*x + q1 = 0
	D = p1*p1 - 4*q1;
	if(D < 0.0)
	{
		retval[0].real( -p1 * 0.5 );
		retval[0].imag( sqrt(-D) * 0.5 );
		retval[1] = conj(retval[0]);
	}
	else
	{
		sqD = sqrt(D);
		retval[0].real( (-p1 + sqD) * 0.5 );
		retval[1].real( (-p1 - sqD) * 0.5 );
	}

	// solving quadratic eq. - x^2 + p2*x + q2 = 0
	D = p2*p2 - 4*q2;
	if(D < 0.0)
	{
		retval[2].real( -p2 * 0.5 );
		retval[2].imag( sqrt(-D) * 0.5 );
		retval[3] = conj(retval[2]);
	}
	else
	{
		sqD = sqrt(D);
		retval[2].real( (-p2 + sqD) * 0.5 );
		retval[3].real( (-p2 - sqD) * 0.5 );
	}

    vector<pair<double, double>> quartics;

    for (int i = 0; i < 4; i++) quartics.push_back(make_pair(retval[i].real(), retval[i].imag()));


    return quartics;
}

unsigned int ToDetection::CalQubic(double *x,double a,double b,double c) 
{
	double a2 = a*a;
    	double q  = (a2 - 3*b)/9;
	double r  = (a*(2*a2-9*b) + 27*c)/54;
    	double r2 = r*r;
	double q3 = q*q*q;
	double A,B;
    	if(r2<q3)
    	{
    		double t=r/sqrt(q3);
    		if( t<-1) t=-1;
    		if( t> 1) t= 1;
    		t=acos(t);
    		a/=3; q=-2*sqrt(q);
    		x[0]=q*cos(t/3)-a;
    		x[1]=q*cos((t+M_2PI)/3)-a;
    		x[2]=q*cos((t-M_2PI)/3)-a;
    		return 3;
    	}
    	else
    	{
    		A =-pow(fabs(r)+sqrt(r2-q3),1./3);
    		if( r<0 ) A=-A;
    		B = (0==A ? 0 : q/A);

		a/=3;
		x[0] =(A+B)-a;
		x[1] =-0.5*(A+B)-a;
		x[2] = 0.5*sqrt(3.)*(A-B);
		if(fabs(x[2])<eps) { x[2]=x[1]; return 2; }

		return 1;
        }
}